'''
比较两个文件夹
/mnt/nas/datasets/diction/coat0908_img_clo_diff
/mnt/nas/datasets/diction/coat0908_img_clo_diff_dino3
求解 两个文件夹图片文件的 差集 A-B  B-A

方式：
    读取 dir/names.txt 获得 所有的 filenames 
    筛选 jpg 后缀的文件 
    将A-B的图片放到 ./A_B
    将B-A的图片放到 ./B_A

    统计 交集数量 差集数量 并 输出
'''
import os
import shutil

def read_jpg_names_from_txt(txt_path):
    if not os.path.exists(txt_path):
        print(f"File not found: {txt_path}")
        return set()
    with open(txt_path, encoding='utf-8') as f:
        names = [line.strip() for line in f if line.strip().endswith('.jpg')]
    return set(names)

def copy_imgs(img_names, src_dir, dst_dir):
    os.makedirs(dst_dir, exist_ok=True)
    for name in img_names:
        src_path = os.path.join(src_dir, name)
        dst_path = os.path.join(dst_dir, name)
        if os.path.exists(src_path):
            shutil.copy2(src_path, dst_path)
        else:
            print(f"Warning: {src_path} not found.")

def main():
    dirA = '/mnt/nas/datasets/diction/coat0908_img_clo_diff'
    dirB = '/mnt/nas/datasets/diction/coat0908_img_clo_diff_dino3'
    namesA_path = os.path.join(dirA, 'names.txt')
    namesB_path = os.path.join(dirB, 'names.txt')

    namesA = read_jpg_names_from_txt(namesA_path)
    namesB = read_jpg_names_from_txt(namesB_path)

    A_B = namesA - namesB
    B_A = namesB - namesA
    inter = namesA & namesB

    print(f"A count: {len(namesA)}")
    print(f"B count: {len(namesB)}")
    print(f"Intersection count: {len(inter)}")
    print(f"A - B count: {len(A_B)}")
    print(f"B - A count: {len(B_A)}")

    copy_imgs(A_B, dirA, './A_B')
    copy_imgs(B_A, dirB, './B_A')

    print("Done.")

if __name__ == '__main__':
    main()
